import torch
from site_packages.ml_libs.nlp.base import Dictionary
from site_packages.utils.job import DataOp
from utils import InputFormatter
from settings import BASE_DIR
import os
import jieba
from jieba import posseg
import re
import numpy as np
from numpy.random import randn


TRAIN_PATH = BASE_DIR + '/data/'
file_train = BASE_DIR + 'train_data.data'
START_TAG = "SOS"
STOP_TAG = "EOS"
# 7个标签
chunk_tags = ['O', 'B-PER', 'I-PER', 'B-LOC', 'I-LOC', "B-ORG", "I-ORG"]


def load_data(train=True):
    """
    :param train: True，训练数据；False，测试数据
    :return: 文本列表和标签列表
    """
    if train:
        file = os.path.join(TRAIN_PATH ,'train_data.data')
    else:
        file = os.path.join(TRAIN_PATH, 'test_data.data')

    with open(file, 'rb') as fp:
        data = fp.read().decode('utf8').split(":")
    print("文本量为：", len(data))
    texts = []
    tags = []
    for line in data:
        contents = line.split()
        text = contents[::2]
        tag = contents[1::2]
        texts.append(text)
        tags.append(tag)
    return texts, tags


def posseg_cut(text):
    tags = posseg.cut(text)
    text_post = []
    tag_post = []
    for tag in tags:
        # print(tag)
        # print(tag.word)
        text_post.append(tag.word)
        tag_post.append(tag.flag)
    return text_post, tag_post


# ------特征抽取部分------
"""
以下用'我们 今天 中午 去 港丽餐厅 吃饭'这句分词后的句子为例进行解释
1、特征矩阵S的形式，此时没有考虑特征系数
词（时间步） 标记 | 转移特征1 | 转移特征2 | 转移特征3 | 状态特征1 | 状态特征2 | ... | 状态特征N | 
我们         r        1         1         0        ...        ...
今天         d        0         0         1        ...        ...
下午         d        1         1         0        ...        ...
去           v        0         0         0        ...       ...
港丽餐厅      n        0         0         0        ...        ...
吃饭         v        1         1         1        ...        ...

状态特征提取模式有两种：
（1）在BiLSTM-CRF中，状态特征就是BiLSTM提取的特征，比如上面的例子 BiLSTM提取了N维的特征
（2）设置人工特征，提取方式同转移特征的方式

2、每个特征有权重

3、
"""


def trans_vn(y_prev, y_cur):
    if y_prev == 'v' and y_cur == 'n':
        return 1.
    else:
        return 0.


def trans_nv(y_prev, y_cur):
    if y_prev == 'n' and y_cur == 'v':
        return 1.
    else:
        return 0.


def get_trans(X, y):
    transition_funcs = [
        trans_vn,
        trans_nv
    ]
    K = len(transition_funcs)
    coefs = [randn() for k in range(K)]
    scores = []
    for k, func in enumerate(transition_funcs):
        score = list(map(lambda prev, cur:func(prev, cur), y[:-1], y[1:]))
        score.append(0)
        score = coefs[k] * np.array(score)
        scores.append(score)

    print(scores)
    print("len(scores):", len(scores))
    return np.stack(scores).T


def get_status(X,y):
    T = len(X)
    dim = 200
    return randn(T, dim)


def log_sum_exp(vec):
    # 减去max的目的避免内存上溢
    # return torch.log(torch.sum(torch.exp(vec)))
    max_score = vec[0, torch.argmax(vec)]
    max_score_broadcast = max_score.view(1,-1).expand(1, vec.size()[1])
    return max_score + \
        torch.log(torch.sum(torch.exp(vec - max_score_broadcast)))


if __name__ == "__main__":
    # texts, tags = load_data()
    # sentences = """新冠肺炎 疫情 正在 全球 多国 反弹，但 美国 政客 却 依然 没有 把 心思 放在 控制 疫情 蔓延 、 守护 民众 生命 健康 上。
    # 美国 继续 炮制 谎言，不断 甩锅 推责，破坏 全球 抗疫 合作。
    # 美国 情报 部门 近日 发布 所谓 新冠病毒 溯源 问题 调查 报告、白宫 发表 有关 声明，罔顾 科学 和 事实，对 中国 进行 污蔑 攻击，其 目的 就是 借 溯源 问题 向 中国 甩锅 推责、散播 政治 病毒。
    # 这 本身 就是 将 溯源 问题 政治化 的 铁证。
    # """
    # sentences = re.sub(" ", "", sentences).split("\n")
    # s_cuts = []
    # tags = []
    # for s in sentences:
    #     s_cut, tag = posseg_cut(s)
    #     s_cuts.append(s_cut)
    #     tags.append(tag)
    #
    # # 使用一句话来执行一次前向传播
    # trans_scores = get_trans(s_cuts[0], tags[0])
    # print("trans_scores.shape:", trans_scores.shape)
    # status_scores = get_status(s_cuts[0], tags[0])
    # print("status_scores.shape:", status_scores.shape)

    # torch测试cat
    # x = torch.arange(10, dtype=torch.float).reshape(1, -1)
    # print(x)
    # # r = log_sum_exp(x)
    # # print(r)
    # y = torch.arange(10, dtype=torch.float).reshape(1, -1) + 2
    # r = torch.cat([x,y], dim=0)
    # print(r)


    # numpy拼接测试
    """
    对于concatenate来说
    axis=0，指的是按照第0维（行）来拼接，此时列数不变，行数增加
    
    对于sum、max等聚合函数来说
    axis=0，指的是沿着第0维（行）求值，此时列数不变，行减少
    总结一下，axis=0，第0维的信息会发生变化
    """
    x = np.arange(10).reshape(1, -1)
    # print(x)
    # # r = log_sum_exp(x)
    # # print(r)
    y = np.arange(10).reshape(1, -1) + 2
    r = np.concatenate([x,y], axis=0)
    # print(r)

    r2 = np.max(r, axis=0)
    print(r2)